import torch
from torch.utils.data import DataLoader

# 设置设备
device = torch.device("mps")

# 初始化模型和数据
model = torch.nn.Linear(20, 1).to(device)
dataset = torch.randn(1000, 20).to(device)
loader = DataLoader(dataset, batch_size=128)

# 训练循环
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
model.train()

for epoch in range(10):
    for batch in loader:
        optimizer.zero_grad()
        output = model(batch)
        loss = output.sum()
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item()}")